#include "Vocab.h"
#include <iostream>

Vocab::Vocab(const std::string& file_path) {
    load(file_path);
}

void Vocab::load(const std::string& file_path) {
    std::ifstream file(file_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open vocabulary file: " + file_path);
    }
    json vocab_data;
    file >> vocab_data;
    token_to_idx = vocab_data["token_to_idx"].get<std::unordered_map<std::string, int>>();
    idx_to_token = vocab_data["idx_to_token"].get<std::vector<std::string>>();
}

int Vocab::operator[](const std::string& token) const {
    auto it = token_to_idx.find(token);
    return (it != token_to_idx.end()) ? it->second : 0;
}

std::string Vocab::to_tokens(int idx) const {
    if (idx >= 0 && idx < idx_to_token.size()) {
        return idx_to_token[idx];
    }
    return "<unk>";
}